ci: Add workflow to build official docker images with multiarch (#83437) Resolves https://github.com/pytorch/pytorch/issues/80764 Signed-off-by: Eli Uriegas <seemethere101@gmail.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/83437 Approved by: https://github.com/ZainRizvi, https://github.com/malfet
diff --git a/Dockerfile b/Dockerfile index 1bd522a..815a910 100644 --- a/Dockerfile +++ b/Dockerfile
@@ -11,8 +11,7 @@ ARG PYTHON_VERSION=3.8 FROM ${BASE_IMAGE} as dev-base -RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \ - apt-get update && apt-get install -y --no-install-recommends \ +RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential \ ca-certificates \ ccache \ @@ -28,9 +27,16 @@ FROM dev-base as conda ARG PYTHON_VERSION=3.8 +# Automatically set by buildx +ARG TARGETPLATFORM +# translating Docker's TARGETPLATFORM into miniconda arches +RUN case ${TARGETPLATFORM} in \ + "linux/arm64") MINICONDA_ARCH=aarch64 ;; \ + *) MINICONDA_ARCH=x86_64 ;; \ + esac && \ + curl -fsSL -v -o ~/miniconda.sh -O "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh" COPY requirements.txt . -RUN curl -fsSL -v -o ~/miniconda.sh -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \ - chmod +x ~/miniconda.sh && \ +RUN chmod +x ~/miniconda.sh && \ ~/miniconda.sh -b -p /opt/conda && \ rm ~/miniconda.sh && \ /opt/conda/bin/conda install -y python=${PYTHON_VERSION} cmake conda-build pyyaml numpy ipython && \ @@ -57,15 +63,21 @@ ARG CUDA_CHANNEL=nvidia ARG INSTALL_CHANNEL=pytorch-nightly ENV CONDA_OVERRIDE_CUDA=${CUDA_VERSION} -RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -c "${CUDA_CHANNEL}" -y python=${PYTHON_VERSION} pytorch torchvision torchtext "cudatoolkit=${CUDA_VERSION}" && \ +# Automatically set by buildx +RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION} +ARG TARGETPLATFORM +# On arm64 we can only install wheel packages +RUN case ${TARGETPLATFORM} in \ + "linux/arm64") pip install --extra-index-url https://download.pytorch.org/whl/cpu/ torch torchvision torchtext ;; \ + *) /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -c "${CUDA_CHANNEL}" -y "python=${PYTHON_VERSION}" pytorch torchvision torchtext "cudatoolkit=${CUDA_VERSION}" ;; \ + esac && \ /opt/conda/bin/conda clean -ya RUN /opt/conda/bin/pip install torchelastic FROM ${BASE_IMAGE} as official ARG PYTORCH_VERSION LABEL com.nvidia.volumes.needed="nvidia_driver" -RUN --mount=type=cache,id=apt-final,target=/var/cache/apt \ - apt-get update && apt-get install -y --no-install-recommends \ +RUN apt-get update && apt-get install -y --no-install-recommends \ ca-certificates \ libjpeg-dev \ libpng-dev && \ 